package com.khalil.kkadmin.module.ai.agent.chatsession;

import dev.langchain4j.data.message.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;
import java.util.Optional;

@Slf4j
@Component
public class JpaChatMemoryStore implements ChatMemoryStore {

    private final ChatSessionRepository chatSessionRepository;

    public JpaChatMemoryStore(ChatSessionRepository chatSessionRepository) {
        this.chatSessionRepository = chatSessionRepository;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        log.info("Getting messages for memoryId: {}", memoryId);
        ChatSession chatSession = chatSessionRepository.findById((Long) memoryId).orElse(null);
        if (chatSession == null) {
            return Collections.emptyList();
        }
        return ChatMessageDeserializer.messagesFromJson(chatSession.getMemoryText());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> list) {
        log.info("Getting messages for memoryId: {}", memoryId);
        Optional<ChatSession> chatSessionOptional = chatSessionRepository.findById((Long) memoryId);
        if (chatSessionOptional.isEmpty()) {
            log.warn("ChatSession not found");
            return;
        }
        ChatSession chatSession = chatSessionOptional.get();
        chatSession.setMemoryText(ChatMessageSerializer.messagesToJson(list));
        chatSessionRepository.save(chatSession);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        log.info("Getting messages for memoryId: {}", memoryId);
        chatSessionRepository.deleteById((Long) memoryId);
    }
}
